package indigo.mybatis.pagination;

import java.util.List;
import java.util.Map;
import java.util.Properties;

import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import com.google.common.collect.Maps;

import indigo.mybatis.MappedStatementTemplate;
import indigo.mybatis.SqlSourceWrapper;
import indigo.util.ImmutablePagedList;

@Intercepts({
		@Signature(type = Executor.class, method = "query", args = {
				MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class
		})
})
public class PaginationInterceptor implements Interceptor {
	private static final String OFFSET_KEY = "_offset";
	private static final String LIMIT_KEY = "_limit";

	private static final String SELECT_COUNT_SUFFIX = "!selectCount";
	private static final String SELECT_PAGE_SUFFIX = "!selectPage";

	private int extractCountResult(List<?> results) {
		Number numberResult = null;
		for (Object result : results) {
			if (numberResult == null) {
				numberResult = (Number) result;
				break;
			}
			throw new IllegalStateException("Count query returns multiple results");
		}
		if (numberResult == null) {
			throw new IllegalStateException("Count query returns nothing");
		}
		return numberResult.intValue();
	}

	@Override
	public Object intercept(Invocation invocation) throws Throwable {
		Object[] args = invocation.getArgs();

		MappedStatement statement = (MappedStatement) args[0];
		RowBounds rowBounds = (RowBounds) args[2];
		ResultHandler<?> resultHandler = (ResultHandler<?>) args[3];

		if (rowBounds instanceof PagedRowBounds) {
			Configuration configuration = statement.getConfiguration();

			String statementId = statement.getId();
			MappedStatement countStatement = configuration.getMappedStatement(statementId + SELECT_COUNT_SUFFIX);
			MappedStatement pageStatement = configuration.getMappedStatement(statementId + SELECT_PAGE_SUFFIX);

			Map<String, Object> extraParams = Maps.newHashMap();
			extraParams.put(OFFSET_KEY, rowBounds.getOffset());
			extraParams.put(LIMIT_KEY, rowBounds.getLimit());

			MappedStatementTemplate countStatementTemplate = new MappedStatementTemplate(countStatement);
			countStatementTemplate.setSqlSource(SqlSourceWrapper.wrap(countStatement.getSqlSource(), extraParams));

			MappedStatementTemplate pageStatementTemplate = new MappedStatementTemplate(pageStatement);
			pageStatementTemplate.setSqlSource(SqlSourceWrapper.wrap(pageStatement.getSqlSource(), extraParams));

			args[2] = RowBounds.DEFAULT;
			args[3] = Executor.NO_RESULT_HANDLER;
			args[0] = countStatementTemplate.create();
			List<?> countResult = (List<?>) invocation.proceed();

			args[3] = resultHandler;
			args[0] = pageStatementTemplate.create();
			List<?> pageResult = (List<?>) invocation.proceed();

			PagedRowBounds pagedRowBounds = (PagedRowBounds) rowBounds;
			return ImmutablePagedList.of(pagedRowBounds.getPageIndex(), pagedRowBounds.getPageSize(),
					extractCountResult(countResult), pageResult);
		} else {
			return invocation.proceed();
		}
	}

	@Override
	public Object plugin(Object target) {
		if (target instanceof Executor) {
			return Plugin.wrap(target, this);
		}
		return target;
	}

	@Override
	public void setProperties(Properties properties) {
	}
}
